# encoding=utf-8

import torch
from seg_common.PredictService import PredictFactory
from seg_common import CommonAnnotation
from PIL import Image
from torchvision import transforms
from seg_system import ApplicationConfig
import numpy as np
import cv2 as cv
import os
from seg_system.segmentation.service.nerveformer.networks.whole_network import Mymodel


class NerveFormerService(PredictFactory):

    @CommonAnnotation.override()
    def before(self, obj, **kwargs):
        mode = kwargs.get('mode', ApplicationConfig.SystemConfig.MODE)

        if mode == ApplicationConfig.SystemConfig.MODE_FOLDER:
            return obj
        elif mode == ApplicationConfig.SystemConfig.MODE_PICTURE:
            image = Image.open(obj)
            gray = image.convert('L')
            trans = self.transform(gray)
            trans = trans.unsqueeze(0)
            trans = trans.to(ApplicationConfig.SystemConfig.DEVICE)
            return trans

    @CommonAnnotation.override()
    def forward(self, obj, **kwargs):
        with torch.no_grad():
            output = self.model(obj)
            return output

    @CommonAnnotation.override()
    def after(self, obj, **kwargs):
        o = obj.detach().cpu().numpy()
        l = []
        for each_o in o:
            each_o = each_o.squeeze(0)
            each_o = (each_o * 255).astype(np.uint8)
            l.append(each_o)
        return l

    def to_list(self, obj):
        if not isinstance(obj, list):
            return [obj]
        return obj

    @CommonAnnotation.override()
    def save(self, obj, path, name, **kwargs):
        obj = self.to_list(obj)
        name = self.to_list(name)

        each_from = kwargs.get('each_from', '')
        tmp_path = kwargs.get('tmp_path')
        assert each_from is not None
        alpha = kwargs.get('alpha', 0.5)
        beta = 1.0 - alpha
        color = kwargs.get('color', np.array([0, 212, 255]))  # yellow bgr

        for each_o, each_n in zip(obj, name):
            file = os.path.join(path, each_n)
            tmp_file = os.path.join(tmp_path, each_n)

            ori_file = cv.imread(os.path.join(each_from, each_n))
            no_zero_index = np.nonzero(each_o)
            ori_file[no_zero_index[0], no_zero_index[1], :] = \
                ori_file[no_zero_index[0], no_zero_index[1], :] * beta + alpha * color

            cv.imwrite(tmp_file, ori_file)
            cv.imwrite(file, each_o)
            assert os.path.exists(file)

    def __init__(self, model_path: str = ""):
        self.model = Mymodel()

        model_path = os.path.join(model_path, ApplicationConfig.NetConfig.SEGMENTATION_DICT[
            ApplicationConfig.NetConfig.SEGMENTATION_USE_NERVE])

        # 继承体系适配
        if model_path != "":
            model_d = torch.load(model_path)
            self.model.load_state_dict(model_d)

        self.model.to(ApplicationConfig.SystemConfig.DEVICE)
        self.model.eval()
        self.transform = transforms.Compose([
            transforms.Resize([384, 384]),
            transforms.ToTensor()
        ])

    def __new__(cls, *args, **kwargs):
        """实际上就是一个标准流程工具类，所以声明成一个单例
        """
        if not hasattr(cls, '_inst'):
            cls._inst = super(NerveFormerService, cls).__new__(cls)
        return cls._inst
